#include "stdafx.h"
#include "Direct3D9Hijack.h"
#include "../../Common/Interfaces/IEnvironment.h"
#include "../../Common/Interfaces/IUserInterface.h"
#include "../../Common/PluginManager/PluginManager.h"
#include "../../Common/Exception.h"

bool gb_HasRenderedUi = false;
bool gb_IsReleasingDevice = false;
bool gb_FirstDevice = false;
IServiceLocator* k_Direct3D9Hijack::smr_ServiceLocator_ = NULL;

typeCreateDevice Real_CreateDevice = NULL;
DETOUR_TRAMPOLINE(LPDIRECT3D9 WINAPI Real_Direct3DCreate9(unsigned int SDKVersion), Direct3DCreate9);
typeReset Real_Reset = NULL;
typeEndScene Real_EndScene = NULL;
typeSetFVF Real_SetFVF = NULL;
typeSetVertexDeclaration Real_SetVertexDeclaration = NULL;
typeRelease Real_Release = NULL;



k_Direct3D9Hijack::k_Direct3D9Hijack()
	: mb_Initialized(false)
	, mb_HasDeviceCreated(false)
{
}


k_Direct3D9Hijack::~k_Direct3D9Hijack()
{
}


void* k_Direct3D9Hijack::QueryInterface(const GUID& ar_InterfaceGuid)
{
	if (ar_InterfaceGuid == __uuidof(IDirect3DHijack)) 
		return static_cast<IDirect3DHijack*>(this);
	else if (ar_InterfaceGuid == __uuidof(IService)) 
		return static_cast<IService*>(this);

	return NULL; 
}


void k_Direct3D9Hijack::Initialize(IServiceLocator& ar_ServiceLocator)
{
	smr_ServiceLocator_ = &ar_ServiceLocator;
	mb_Initialized = true;

	unsigned int uiReturn = DetourFunctionWithTrampoline((unsigned char*)Real_Direct3DCreate9, 
		(unsigned char*)k_Direct3D9Hijack::Direct3DCreate9Replace);
	LPDIRECT3D9 lr_D3D9_ = Direct3DCreate9(D3D_SDK_VERSION);
	if (lr_D3D9_ != NULL)
		IDirect3D9_Release(lr_D3D9_);
}


void k_Direct3D9Hijack::Deinitialize()
{
	if (Real_CreateDevice != NULL)
		DetourRemove((PBYTE)Real_CreateDevice, (PBYTE)k_Direct3D9Hijack::CreateDeviceReplace);
	Real_CreateDevice = NULL;
	if (Real_Release != NULL)
		DetourRemove((PBYTE)Real_Release, (PBYTE)k_Direct3D9Hijack::ReleaseReplace);
	Real_Release = NULL;
	if (Real_Reset != NULL)
		DetourRemove((PBYTE)Real_Reset, (PBYTE)k_Direct3D9Hijack::ResetReplace);
	Real_Reset = NULL;
	if (Real_EndScene != NULL)
		DetourRemove((PBYTE)Real_EndScene, (PBYTE)k_Direct3D9Hijack::EndSceneReplace);
	Real_EndScene = NULL;
	if (Real_SetFVF != NULL)
		DetourRemove((PBYTE)Real_SetFVF, (PBYTE)k_Direct3D9Hijack::SetFVFReplace);
	Real_SetFVF = NULL;
	if (Real_SetVertexDeclaration != NULL)
		DetourRemove((PBYTE)Real_SetVertexDeclaration, (PBYTE)k_Direct3D9Hijack::SetVertexDeclarationReplace);
	Real_SetVertexDeclaration = NULL;

	mb_Initialized = false;
}


bool k_Direct3D9Hijack::get_IsInitialized() const
{
	return mb_Initialized;
}


bool k_Direct3D9Hijack::get_HasDeviceCreated() const
{
	return mb_HasDeviceCreated;
}


void k_Direct3D9Hijack::set_HasDeviceCreated(bool ab_State)
{
	mb_HasDeviceCreated = ab_State;
}


HRESULT k_Direct3D9Hijack::CallEndScene(IDirect3DDevice9* ar_Device_)
{
	return Real_EndScene(ar_Device_);
}



////////////////////////////////////////////////////////////////////////////////////////////
// Hijacked function calls
////////////////////////////////////////////////////////////////////////////////////////////
HRESULT WINAPI k_Direct3D9Hijack::CreateDeviceReplace(typeCreateDeviceParam)
{
	IServiceLocator& lr_SvLocator = k_Direct3D9Hijack::get_ServiceLocator();
	IPluginManager& lr_PluginManager = object_cast<IPluginManager>(lr_SvLocator);
	IDirect3DHijack& lr_D3DHijack = object_cast<IDirect3DHijack>(lr_SvLocator);

	IEnvironment& lr_Environment = object_cast<IEnvironment>(lr_SvLocator);
	std::wstring ls_LogFile = lr_Environment.get_ModulePath()+L"Walhall.log";

	if (!gb_FirstDevice && !lr_D3DHijack.get_HasDeviceCreated())
	{
		gb_FirstDevice = true;
		
		NewLog(ls_LogFile);
		PrintLog(ls_LogFile, L"Walhall: Device was successfully build up.\n");
		PrintLog(ls_LogFile, L"Walhall: Modulepath: " << lr_Environment.get_ModulePath() << L"\n");
		PrintLog(ls_LogFile, L"Walhall: Workpath: " << lr_Environment.get_WorkPath() << L"\n");

		lr_PluginManager.LoadPlugins();
	}

	//SE: force non pure device!
	BehaviorFlags &= (~D3DCREATE_PUREDEVICE);

	HRESULT hr = Real_CreateDevice(lpd3d,Adapter,DeviceType,hFocusWindow, BehaviorFlags,pPresentationParameters,ppReturnedDeviceInterface);

	if (!lr_D3DHijack.get_HasDeviceCreated())
	{
		PrintLog(ls_LogFile, L"Walhall: Creating device.\n");

		// call plugin functions
		std::vector<IDirect3D9BasePlugin*> lk_Plugins = lr_PluginManager.get_Plugins();
		for(unsigned int i=0; i<lk_Plugins.size(); i++)
		{
			try
			{
				if (lk_Plugins[i]->IsInitialized())
					lk_Plugins[i]->Deinitialize();
			}
			catch (Exception& e)
			{
				PrintLog(ls_LogFile, L"Fatal Error: Deinitialize of " << lk_Plugins[i]->get_Name() << L" failed. \n");
				PrintLog(ls_LogFile, e.GetMessage() << L"\n");
			}

			try
			{
				lk_Plugins[i]->Initialize(*ppReturnedDeviceInterface);
			}
			catch (Exception& e)
			{
				PrintLog(ls_LogFile, L"Fatal Error: Initialize of " << lk_Plugins[i]->get_Name() << L" failed. \n");
				PrintLog(ls_LogFile, e.GetMessage() << L"\n");
			}
		}

		HIJACKVTABLEFUNCTION(ls_LogFile, L"Core DLL", typeRelease, Release, *ppReturnedDeviceInterface, k_Direct3D9Hijack);
		HIJACKVTABLEFUNCTION(ls_LogFile, L"Core DLL", typeReset, Reset, *ppReturnedDeviceInterface, k_Direct3D9Hijack);
		HIJACKVTABLEFUNCTION(ls_LogFile, L"Core DLL", typeEndScene, EndScene, *ppReturnedDeviceInterface, k_Direct3D9Hijack);
		//HIJACKVTABLEFUNCTION(ls_LogFile, L"Core DLL", typeSetFVF, SetFVF, *ppReturnedDeviceInterface, k_Direct3D9Hijack);
		//HIJACKVTABLEFUNCTION(ls_LogFile, L"Core DLL", typeSetVertexDeclaration, SetVertexDeclaration, *ppReturnedDeviceInterface, k_Direct3D9Hijack);
		
		IUserInterface& lr_Ui = object_cast<IUserInterface>(lr_SvLocator);
		lr_Ui.set_Device(*ppReturnedDeviceInterface);

		gb_IsReleasingDevice = false;

		lr_D3DHijack.set_HasDeviceCreated(true);

		PrintLog(ls_LogFile, L"Walhall: Creating device finished.\n");
	}

	return hr;
}

ULONG WINAPI k_Direct3D9Hijack::ReleaseReplace(typeReleaseParam)
{
	IServiceLocator& lr_SvLocator = k_Direct3D9Hijack::get_ServiceLocator();
	IEnvironment& lr_Environment = object_cast<IEnvironment>(lr_SvLocator);
	IDirect3DHijack& lr_D3DHijack = object_cast<IDirect3DHijack>(lr_SvLocator);

	std::wstring ls_LogFile = lr_Environment.get_ModulePath()+L"Walhall.log";

	bool lb_Last = lpd3dd->lpVtbl->AddRef(lpd3dd) == 2;
	
	Real_Release(lpd3dd);

	IPluginManager& lr_PluginManager = object_cast<IPluginManager>(lr_SvLocator);

	// call plugin functions
	if (lb_Last && !gb_IsReleasingDevice)
	{
		PrintLog(ls_LogFile, L"Walhall: Releasing device.\n");

		gb_IsReleasingDevice = true;

		std::vector<IDirect3D9BasePlugin*> lk_Plugins = lr_PluginManager.get_Plugins();
		for(unsigned int i=0; i<lk_Plugins.size(); i++)
			try
			{
				if (lk_Plugins[i]->IsInitialized())
					lk_Plugins[i]->Deinitialize();
			}
			catch (Exception& e)
			{
				PrintLog(ls_LogFile, L"Fatal Error: Deinitialize of " << lk_Plugins[i]->get_Name() << L" failed. \n");
				PrintLog(ls_LogFile, e.GetMessage() << L"\n");
			}


		IUserInterface& lr_Ui = object_cast<IUserInterface>(lr_SvLocator);
		lr_Ui.set_Device(NULL);

		lr_D3DHijack.set_HasDeviceCreated(false);

		PrintLog(ls_LogFile, L"Walhall: Releasing device finished.\n");
	}

	return Real_Release(lpd3dd);
}


LPDIRECT3D9 WINAPI k_Direct3D9Hijack::Direct3DCreate9Replace(unsigned int SDKVersion)
{
	IEnvironment& lr_Environment = object_cast<IEnvironment>(k_Direct3D9Hijack::get_ServiceLocator());
	std::wstring ls_LogFile = lr_Environment.get_ModulePath()+L"Walhall.log";
	
	PrintLog(ls_LogFile, L"DLL: DirectX9 created (SDKVersion: " << SDKVersion << L")\n");
	LPDIRECT3D9 pDirect3D = Real_Direct3DCreate9(SDKVersion);

	HIJACKVTABLEFUNCTION(ls_LogFile, L"Core DLL", typeCreateDevice, CreateDevice, pDirect3D, k_Direct3D9Hijack);

	DetourRemove((unsigned char*)Real_Direct3DCreate9, (unsigned char*)Direct3DCreate9Replace);

	return pDirect3D;
}


HRESULT WINAPI k_Direct3D9Hijack::ResetReplace(typeResetParam)
{
	IServiceLocator& lr_SvLocator = k_Direct3D9Hijack::get_ServiceLocator();
	IPluginManager& lr_PluginManager = object_cast<IPluginManager>(lr_SvLocator);
	IEnvironment& lr_Environment = object_cast<IEnvironment>(lr_SvLocator);

	std::wstring ls_LogFile = lr_Environment.get_ModulePath()+L"Walhall.log";

	PrintLog(ls_LogFile, L"DLL: Reset start...\n");

	// call plugin functions
	std::vector<IDirect3D9BasePlugin*> lk_Plugins = lr_PluginManager.get_Plugins();
	for(unsigned int i=0; i<lk_Plugins.size(); i++)
	{
		try
		{
			if (lk_Plugins[i]->IsInitialized())
				lk_Plugins[i]->OnBeforeReset();
		}
		catch (Exception& e)
		{
			PrintLog(ls_LogFile, L"Fatal Error: OnBeforeReset of " << lk_Plugins[i]->get_Name() << L" failed. \n");
			PrintLog(ls_LogFile, e.GetMessage() << L"\n");
		}
	}

	IUserInterface& lr_Ui = object_cast<IUserInterface>(lr_SvLocator);
	
	lr_Ui.set_Device(NULL);
	HRESULT result = Real_Reset(lpd3dd,pPresentationParameters);
	lr_Ui.set_Device(lpd3dd);

	// call plugin functions
	for(unsigned int i=0; i<lk_Plugins.size(); i++)
	{
		try
		{
			if (lk_Plugins[i]->IsInitialized())
				lk_Plugins[i]->OnAfterReset();
		}
		catch (Exception& e)
		{
			PrintLog(ls_LogFile, L"Fatal Error: OnAfterReset of " << lk_Plugins[i]->get_Name() << L" failed. \n");
			PrintLog(ls_LogFile, e.GetMessage() << L"\n");
		}
	}

	PrintLog(ls_LogFile, L"DLL: Reset end...\n");

	return result;
}


HRESULT WINAPI k_Direct3D9Hijack::EndSceneReplace(typeEndSceneParam)
{
	if (!gb_HasRenderedUi)
		RenderWalhall(lpd3dd);

	HRESULT result = Real_EndScene(lpd3dd);
	gb_HasRenderedUi = false;

	return result;
}


HRESULT WINAPI k_Direct3D9Hijack::SetFVFReplace(typeSetFVFParam)
{
	IServiceLocator& lr_SvLocator = k_Direct3D9Hijack::get_ServiceLocator();
	IPluginManager& lr_PluginManager = object_cast<IPluginManager>(lr_SvLocator);
	IEnvironment& lr_Environment = object_cast<IEnvironment>(lr_SvLocator);

	std::wstring ls_LogFile = lr_Environment.get_ModulePath()+L"Walhall.log";
	if ((FVF & D3DFVF_XYZRHW) == D3DFVF_XYZRHW && !gb_HasRenderedUi)
	{
		PrintLog(ls_LogFile, L"DLL: Render through SetFVF\n");
		RenderWalhall(lpd3dd);
	}
	HRESULT result = Real_SetFVF(lpd3dd, FVF);

	return result;
}


HRESULT WINAPI k_Direct3D9Hijack::SetVertexDeclarationReplace(typeSetVertexDeclarationParam)
{
	IServiceLocator& lr_SvLocator = k_Direct3D9Hijack::get_ServiceLocator();
	IPluginManager& lr_PluginManager = object_cast<IPluginManager>(lr_SvLocator);
	IEnvironment& lr_Environment = object_cast<IEnvironment>(lr_SvLocator);

	std::wstring ls_LogFile = lr_Environment.get_ModulePath()+L"Walhall.log";
	if (!gb_HasRenderedUi)
	{
		unsigned int lui_ElementCount = 0;
		HRESULT hr = IDirect3DVertexDeclaration9_GetDeclaration(pDecl, NULL, &lui_ElementCount);
		D3DVERTEXELEMENT9* lr_Decl_ = new D3DVERTEXELEMENT9[lui_ElementCount];
		hr = IDirect3DVertexDeclaration9_GetDeclaration(pDecl, lr_Decl_, &lui_ElementCount);
		for (unsigned int i = 0; i < lui_ElementCount; ++i)
		{
			if (lr_Decl_[i].Usage == D3DDECLUSAGE_POSITIONT)
			{
				PrintLog(ls_LogFile, L"DLL: Render through SetVertexDeclaration\n");
				RenderWalhall(lpd3dd);
				break;
			}
		}
		delete [] lr_Decl_;
	}
	HRESULT result = Real_SetVertexDeclaration(lpd3dd, pDecl);

	return result;
}


void k_Direct3D9Hijack::RenderWalhall(IDirect3DDevice9* lpd3dd)
{
	gb_HasRenderedUi = true;

	IServiceLocator& lr_SvLocator = k_Direct3D9Hijack::get_ServiceLocator();
	IPluginManager& lr_PluginManager = object_cast<IPluginManager>(lr_SvLocator);

	IEnvironment& lr_Environment = object_cast<IEnvironment>(lr_SvLocator);
	std::wstring ls_LogFile = lr_Environment.get_ModulePath()+L"Walhall.log";

	// call plugin functions
	std::vector<IDirect3D9BasePlugin*> lk_Plugins = lr_PluginManager.get_Plugins();
	for(unsigned int i=0; i<lk_Plugins.size(); i++)
	{
		try
		{
			if (!lk_Plugins[i]->IsInitialized())
				lk_Plugins[i]->Initialize(lpd3dd);
		}
		catch (Exception& e)
		{
			PrintLog(ls_LogFile, L"Fatal Error: Initialize of " << lk_Plugins[i]->get_Name() << L" failed. \n");
			PrintLog(ls_LogFile, e.GetMessage() << L"\n");
		}

		try
		{
			if (lk_Plugins[i]->IsInitialized() && lk_Plugins[i]->get_Active())
				lk_Plugins[i]->Render();
		}
		catch (Exception& e)
		{
			PrintLog(ls_LogFile, L"Fatal Error: Render of " << lk_Plugins[i]->get_Name() << L" failed. \n");
			PrintLog(ls_LogFile, e.GetMessage() << L"\n");
		}

	}
	
	IUserInterface& lr_Ui = object_cast<IUserInterface>(lr_SvLocator);
	lr_Ui.RenderMenu();

}